import numpy as np
from scipy.signal import butter, lfilter
from scipy.signal import savgol_filter
from config import sample_frequency

current_processing_arr = [0]


def butter_lowpass(filter_params_dic):
    f_low = filter_params_dic['low'] / 60
    f_hi = filter_params_dic['hi'] / 60
    # print(f'取滤波范围为：{f_low}, {f_hi}')
    low = f_low * 2 / sample_frequency
    hi = f_hi * 2 / sample_frequency
    b, a = butter(1, [low, hi], btype='bandpass', analog=False)
    return b, a


def butter_lowpass_filter(data, filter_params_dic):
    b, a = butter_lowpass(filter_params_dic)
    y = 1 * lfilter(b, a, data) + 550
    ans = np.clip(y, 0, 1100)
    return ans


def smooth_data(data, window_length):
    # return savgol_filter(data, window_length=window_length, polyorder=2)
    return moving_average(data, window_length)


def moving_average(data, window_size):
    weights = np.repeat(1.0, window_size) / window_size
    return np.convolve(data, weights, 'valid')


def do_filter(data, filter_params_dic):
    return butter_lowpass_filter(data, filter_params_dic)
    # return fft_filter(data, filter_params_dic)


def fft_filter(data, filter_params_dic):
    fft_result = np.fft.fft(data)
    freqs = np.fft.fftfreq(len(data), 0.02)
    f_low = filter_params_dic['low'] / 60
    f_hi = filter_params_dic['hi'] / 60
    filtered_fft_result = fft_result.copy()
    filtered_fft_result[np.abs(freqs) < f_low] = 0
    filtered_fft_result[np.abs(freqs) > f_hi] = 0
    filtered_data = np.fft.ifft(filtered_fft_result)
    real_data = np.real(filtered_data)
    clip_data = np.clip(real_data, 0, 1100)
    fit_data = clip_data + 550
    return fit_data
